import argparse
import json
from .io import parse_usv_file, parse_port_file, parse_task_file
from .allocator import assign_tasks
from .visualizer import draw_trajectories


def main():
    parser = argparse.ArgumentParser(description="USV Task Allocator CLI")
    parser.add_argument('--usv', required=True, help="Path to USV info txt")
    parser.add_argument('--port', required=True, help="Path to port info txt")
    parser.add_argument('--task', required=True, help="Path to task info txt")
    parser.add_argument('--out', default='results.json', help="Output JSON file")
    parser.add_argument('--lambda_weight', type=float, default=0.1, help="Weight for load redundancy")
    parser.add_argument('--map', type=str, help="Path to base map image (optional)")
    parser.add_argument('--vis_out', type=str, default='trajectories_map.png', help="Trajectory image output path")

    args = parser.parse_args()

    usvs = parse_usv_file(args.usv)
    ports = parse_port_file(args.port)
    tasks = parse_task_file(args.task)

    assignments = assign_tasks(usvs, tasks, lambda_weight=args.lambda_weight)

    for a in assignments:
        print(f"USV {a['usv_id']} → Task {a['task_id']}, Payload: {a['payload']}, Cost: {a['total_cost']:.2f}")

    # 写入 JSON 文件
    json_result = [{
        "usv_id": a["usv_id"],
        "task_id": a["task_id"],
        "time_sequence": a["time_sequence"]
    } for a in assignments]

    with open(args.out, 'w') as f:
        json.dump(json_result, f, indent=2)

    print(f"\n任务分配已保存至：{args.out}")

    import cv2  # 确保文件顶部导入了 OpenCV

    if args.map:
        map_img = cv2.imread(args.map, cv2.IMREAD_COLOR)
        if map_img is None:
            raise FileNotFoundError(f"地图文件未找到: {args.map}")

        draw_trajectories(
            map_img=map_img,
            usvs=usvs,
            ports=ports,
            assignments=assignments,
            out_path=args.vis_out,
            scale=0.1
        )

